import numpy as np


# 加载数据
def load_data():
    with open('../data/约会数据.txt') as fr:
        lines = fr.readlines()

    x = np.empty((len(lines), 3), dtype=float)
    y = np.empty((len(lines)), dtype=int)

    for i in range(len(lines)):
        # str.strip()：删除字符串前后（左右两侧）的空格或特殊字符。
        line = lines[i].strip().split('\t')
        x[i] = line[:3]
        y[i] = line[3]

    return x, y


# 数据归一化，防止大数据变量的干扰，保证每一列的话语权一样
def norm(x):
    x_col_min = x.min(axis=0)
    x_col_max = x.max(axis=0)

    x_norm = x - x_col_min  # 这一步之后，最小值为0

    x_norm /= x_col_max - x_col_min  # 这一步之后，最大值为1
    return x_norm


x, y = load_data()
x = norm(x)
# 切分数据集
train_x = x[:900]
train_y = y[:900]
test_x = x[900:]
test_y = y[900:]


def knn(_x, k=3):
    temp = _x - train_x
    temp = np.power(temp, 2)
    temp = temp.sum(axis=1)
    temp = np.sqrt(temp)
    argsort = temp.argsort()  # [3 2 1 0]
    result = y[argsort][:k]  # [1 1 0]
    return np.bincount(result).argmax()  # 取众数


# 进行测试
correct = 0
for i in range(len(test_x)):
    pred = knn(test_x[i], k=5)
    if pred == test_y[i]:
        correct += 1
print(correct / len(test_x))
